package com.zck2

import com.zck2.utils.UDFPartitioner
import org.apache.spark.SparkContext
import org.apache.spark.rdd.RDD
import org.apache.spark.storage.StorageLevel

import scala.collection.mutable
import scala.collection.mutable.ArrayBuffer

object CrossMatch extends Serializable {
//将原始rdd自定义分区处理
def udfPartitions(sc: SparkContext, pathDir: String, hpx_level_cha: Int): (RDD[(String, String, String, String, String, String, String)], RDD[(String, String, String, String, String, String, String)]) = {
  val rddInit: RDD[Array[String]] = sc.textFile(pathDir).map(_.split(",")).persist(StorageLevel.MEMORY_AND_DISK_SER)
  //过滤无效星体：星等=99.0
  val rdd: RDD[Array[String]] = rddInit.filter(!(_: Array[String])(6).equals("99.0"))
  //移位前索引 为k  并提前计算arr_k2
  val arr_k2: Array[Long] = rdd.map(arr => arr(3).toLong >> (hpx_level_cha * 2)).distinct().collect()
  val kvRdd_hpx1: RDD[(Long, (String, String, Long, String, String, String))] = rdd.map(arr => (arr(2).toLong >> (hpx_level_cha * 2), (arr(0), arr(1), arr(3).toLong >> (hpx_level_cha * 2), arr(4), arr(5), arr(6))))
  val arr_k1: Array[Long] = kvRdd_hpx1.map(_._1).distinct().collect()
  //获取分区后双层索引下的rdd(移位前的)
  val kvNewRdd_hpx1: RDD[(Long, (String, String, Long, String, String, String))] = kvRdd_hpx1.partitionBy(new UDFPartitioner(arr_k1))
    .mapPartitions((par: Iterator[(Long, (String, String, Long, String, String, String))]) => getNewRDD_Index(par))

  //更换移位后索引 为k  移除repartition(1)
  val kvRdd_hpx2: RDD[(Long, (String, String, Long, String, String, String))] = kvNewRdd_hpx1.map(x => (x._2._3, (x._2._1, x._2._2, x._1, x._2._4, x._2._5, "0")))

  // 首先过滤掉 key == -1 的数据，再进行分区
  val filteredKvRdd_hpx2: RDD[(Long, (String, String, Long, String, String, String))] = kvRdd_hpx2.filter { case (key, _) =>
    key != -1  // 过滤掉 key 为 -1 的数据
  }
  //移位后的重分区
  val kvNewRdd_hpx2: RDD[(Long, (String, String, Long, String, String, String))] = filteredKvRdd_hpx2.partitionBy(new UDFPartitioner(arr_k2))

  //分区后 恢复原始格式
  val paredRDD1: RDD[(String, String, String, String, String, String, String)] = kvNewRdd_hpx1.mapPartitions(par => recoveryFormat1(par))
  val paredRDD2: RDD[(String, String, String, String, String, String, String)] = kvNewRdd_hpx2.mapPartitions(par => recoveryFormat2(par))

  (paredRDD1, paredRDD2)
}
  //初始化 sampleData、MatchData
  def initAB(parIter: Iterator[(String, String, String, String, String, String, String)]): (ArrayBuffer[Array[String]], ArrayBuffer[Array[String]], Array[String]) = {

    val sampleData: ArrayBuffer[Array[String]] = ArrayBuffer[Array[String]]()
    val matchData: ArrayBuffer[Array[String]] = ArrayBuffer[Array[String]]()

    var oneStar: Array[String] = null
    var flag = ""
    // 如果 Iterator 为空，直接返回默认值
    if (!parIter.hasNext) {
      return (sampleData, matchData, Array[String]())  // 如果 Iterator 为空，返回空的 sampleData 和 matchData，以及空的 oneStar
    }
    val firstStar: Array[String] = parIter.next().productIterator.toArray.map { any: Any => any.toString }
    val firstFlag: String = firstStar(0)
    sampleData += firstStar
    var temp_in = 0 //符合当前flag为0  不符合为1 用来跳出循环
    while (parIter.hasNext && temp_in == 0) {
      oneStar = parIter.next().productIterator.toArray.map { any: Any => any.toString }
      if (oneStar(0) == firstFlag) {
        sampleData += oneStar
      }
      else {
        temp_in = 1
        flag = oneStar(0) //更新flag
        matchData += oneStar
      }
    }

    while (parIter.hasNext && temp_in == 1) {
      oneStar = parIter.next().productIterator.toArray.map { any: Any => any.toString }
      if (oneStar(0) == flag) {
        matchData += oneStar
      }
      else {
        temp_in = 0
      }
    }
    (sampleData, matchData, oneStar)
  }

  //  实现 两“天”数据的证认  返回新的 sample 和 match
  def matchAB(sampleData: ArrayBuffer[Array[String]], matchData: ArrayBuffer[Array[String]], hMap: mutable.HashMap[String, String]): mutable.HashMap[String, String] = {
    //证认
    val newData: ArrayBuffer[Array[String]] = ArrayBuffer[Array[String]]()
    for (md <- matchData) {
      var is_ok = 0 // 标记 用于判断md 是否证认成功
      for (sd <- sampleData) {
        val dis: Double = Math.sqrt(Math.pow((sd(4).toDouble - md(4).toDouble) * Math.cos((sd(5).toDouble + md(5).toDouble) / 2.0), 2) + Math.pow(sd(5).toDouble - md(5).toDouble, 2))


        if (dis < 3 * Math.sqrt(2 * Math.pow(1.0 / 360, 2))) { //误差：1角秒
          is_ok = 1 //表示 md与 该sampleData 中的星体 成功证认
          sd(6) = "0" //一旦sd 成功被证认 sd(6) 清零
          var k = ""
          var v = ""

          k = sd(2) + "_" + sd(1)
          v = md(2) + "_" + md(1) //156818333285_0001_1

          //将证认结果存储到Map中
          if (hMap.contains(k)) hMap(k) += ("," + v)
          else hMap += k -> v

        }
        else {
          sd(6) = (sd(6).toInt + 1).toString // 记录星体未证认成功次数
        }

        if (sd(6).equals(matchData.size.toString)) {
          sd(6) = (matchData.size - sd(6).toInt + 1).toString // 改为 星体未证认成功天数
        }
      }
      //md证认失败
      if (is_ok == 0) {
        newData += md
      }
    }
    //更新参考星表
    sampleData ++= newData
    hMap
  }
  //更新matchData
  def updateB(parIter: Iterator[(String, String, String, String, String, String, String)], firstStar: Array[String]): (ArrayBuffer[Array[String]], Array[String]) = {

    val matchData = ArrayBuffer[Array[String]]()

    matchData += firstStar
    var oneStar = Array[String]()

    var temp_in = 0
    while (parIter.hasNext && temp_in == 0) {
      oneStar = parIter.next().productIterator.toArray.map {
        any: Any => any.toString
      }
      if (oneStar(0) == firstStar(0)) {
        matchData += oneStar
      }
      else temp_in = 1
    }
    (matchData, oneStar)

  }

  //  递归证认整个天区的所有时间段 保存到hMap中
  def matchAll(paredRDD: RDD[(String, String, String, String, String, String, String)], bound: Int): RDD[(String, Array[String])] = {
    //    val Rdd = paredRDD.repartition(1) //单分区  测试用
    val pRdd = paredRDD

    val HashMapRDD: RDD[mutable.HashMap[String, String]] = pRdd.mapPartitions(parIter => {
      var hMap = new mutable.HashMap[String, String]()
      var sampleData = ArrayBuffer[Array[String]]()
      var matchData = ArrayBuffer[Array[String]]()
      var oneStar: Array[String] = null

      //获取初始sample 和 matchData
      val init: (ArrayBuffer[Array[String]], ArrayBuffer[Array[String]], Array[String]) = initAB(parIter)
      sampleData = init._1
      matchData = init._2
      oneStar = init._3 //下一"天"的第一条星体

      hMap = matchAB(sampleData, matchData, hMap)
      while (parIter.hasNext) {

        //更新 matchData 及 oneStar
        val ud: (ArrayBuffer[Array[String]], Array[String]) = updateB(parIter, oneStar)
        matchData = ud._1
        oneStar = ud._2
        //移除假体
        sampleData = delFalseStar(sampleData, bound)
        //更新后继续证认
        matchAB(sampleData, matchData, hMap)
      }

      Iterator(hMap) //返回每次每个分区证认完成后的结果 集合
    })

    //    HashMapRDD.foreach(println)
    val HMapRdd = HashMapRDD.filter(_.nonEmpty) //重点  手动去空Map()  --在转存mysql时会自动去空
    val newHMaprdd = formatHMapRDD(HMapRdd) //格式化字段
    val processedRDD = newHMaprdd.map { case (key, array) =>
      // 对数组中的元素按下划线分割后的数字进行排序，提取出数字部分
      val sortedArray = array.sortBy { elem =>
        val parts = elem.split("_")
        parts(1).toInt  // 获取第二部分数字并转换为整数
      }

      // 提取 key 中的数字部分
      val keyParts = key.split("_")
      val keyValue = keyParts(1).toInt

      // 提取 sortedArray 中最小值的数字部分
      val minArrayValue = sortedArray.headOption match {
        case Some(elem) =>
          val arrayValue = elem.split("_")(1).toInt
          arrayValue
        case None => Int.MaxValue  // 如果 array 为空，返回最大值
      }

      // 如果 minArrayValue < keyValue，则交换 key 和 value
      if (minArrayValue < keyValue) {
        (sortedArray.head, Array(key) ++ sortedArray.tail)  // 交换
      } else {
        (key, sortedArray)  // 不交换
      }
    }
    processedRDD
  }
  //星体假体的移除
  def delFalseStar(sampleData: ArrayBuffer[Array[String]], bound: Int): ArrayBuffer[Array[String]] = {
    for (elem <- sampleData) {
      if (elem(6).toInt >= bound && elem(5).toDouble > 18) { //未证认次数>bound 或 mag超过阈值16/18
        sampleData -= elem
      }
    }
    sampleData
  }
  //格式复原
  def recoveryFormat1(par: Iterator[(Long, (String, String, Long, String, String, String))]): Iterator[(String, String, String, String, String, String, String)] = {
    par.map(x => (x._2._1, x._2._2, x._1.toString, x._2._3.toString, x._2._4, x._2._5, "0")) //恢复原始格式
  }

  def recoveryFormat2(par: Iterator[(Long, (String, String, Long, String, String, String))]): Iterator[(String, String, String, String, String, String, String)] = {
    par.map(x => (x._2._1, x._2._2, x._2._3.toString, x._1.toString, x._2._4, x._2._5, "0")) //恢复原始格式
  }
  //获取二层索引对应关系 分区内编号
  def getNewRDD_Index(par: Iterator[(Long, (String, String, Long, String, String, String))]): Iterator[(Long, (String, String, Long, String, String, String))] = {
    var flag = "-1"
    var cnt = 1
    var x22 = ""
    par.map(
      //      (h1,(1,0001_1,h2,101.6994,-74.54577,7.5345))
      (x: (Long, (String, String, Long, String, String, String))) => {
        if (!x._2._1.equals(flag)) {
          flag = x._2._1
          cnt = 1
          x22 = x._2._2 + "-" + cnt
        } else {
          cnt += 1
          x22 = x._2._2 + "-" + cnt
        }
        // (h1,(1,0001_1-1,h2,101.6994,-74.54577,7.5345))
        (x._1, (x._2._1, x22, x._2._3, x._2._4, x._2._5, x._2._6))
      }
    )
  }
  //格式化HMapRDD
  def formatHMapRDD(hashMapRDD: RDD[mutable.HashMap[String, String]]): RDD[(String, Array[String])] = {
    val rddInit = hashMapRDD.flatMap(_.map {
      case (k, v) => {
        val cleanedValues = v.split(",").distinct
        (k, cleanedValues)
      }
    })
    rddInit
  }
  //保存分区后的文件
  def paredRDDSave(paredRDD: RDD[(String, String, String, String, String, String, String)], outPath: String): Unit = {
    paredRDD.saveAsTextFile(outPath)
  }
}
